#include "mex.h"
#include <string.h>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include<iostream>
using namespace std;
#define INF 1e9
#define MAXK 100

inline double dist(const double*a, const double* c, int i, int j)
{
	double x = a[3 * i] - c[3 * j];
	double y = a[3 * i + 1] - c[3 * j + 1];
	double z = a[3 * i + 2] - c[3 * j + 2];
	return sqrt(x*x+y*y+z*z);
}

int findCenter3(const double* a, const double* c, int i, int k)
{
	double min = INF;
	int t = -1;
	for (int j = 0; j < k; j++)
	{
		if ( dist(a, c, i, j) < min )
		{
			min = dist(a, c, i, j);
			t = j;
		}
	}
	return t;
}

void initCast(const double *a, int n, int k0,
	double* C, int*P, int* Count)
{
	int h = n / k0;
	double* c = new double [k0 * 3];
	int* p = new int [n];
	int* count = new int[k0];
	memset(c, 0, sizeof(double)*k0*3);
	memset(p, 0, sizeof(int)*n);
	memset(count, 0, sizeof(int)*k0);
	for (int i = 0; i < k0; i++)
	{
		memcpy(c+3*i, a+3*i*h, sizeof(double)*3);
	}
	for (int i = 0; i < n; i++)
	{
		p[i] = findCenter3(a, c, i, k0);
		count[p[i]]++;
	}
	memset(c, 0, sizeof(double)*3*k0);
	for(int i = 0; i < n; i++)
	{
		c[3*p[i]] += a[3*i];
		c[3*p[i]+1] += a[3*i+1];
		c[3*p[i]+2] += a[3*i+2];
	}
	for (int i = 0; i < k0; i++)
	{
		c[3*i] /= count[i];
		c[3*i+1] /= count[i];
		c[3*i+2] /= count[i];
	}
	memcpy(P, p, sizeof(int)*n);
	memcpy(C, c, sizeof(double)*3*k0);
	memcpy(Count, count, sizeof(int)*k0);
}

void cast3(const double* a, int n,int N0, 
	int* K, double* C, int* P, int* Count, int* err)
{
	while( 1 )
	{
		int k = *K;
		double* c = new double [k * 3];
		memcpy(c, C, sizeof(double)*k*3);
		int* p = new int [n];
		memcpy(p, P, sizeof(int)*n);
		int* count = new int[k];
		memset(count, 0, sizeof(int)*k);
		int variation  = 0;
		for (int i = 0; i < n; i++)
		{
			int t = findCenter3(a, c, i, k);
			variation += t - p[i] > 0 ? t - p[i] : p[i] - t;
			p[i] = t;
			count[p[i]]++;
		}
		//check
		int cntFlags = 0;
		for(int i = 0; i < k; i++)
		{
			if(count[i] > N0) cntFlags++;
		}
		if (cntFlags < k) 
		{
			k--;
		}
		else
		{
			memset(c, 0, sizeof(double)*3*k);
			for(int i = 0; i < n; i++)
			{
				c[3*p[i]] += a[3*i];
				c[3*p[i]+1] += a[3*i+1];
				c[3*p[i]+2] += a[3*i+2];
			}
			for (int i = 0; i < k; i++)
			{
				c[3*i] /= count[i];
				c[3*i+1] /= count[i];
				c[3*i+2] /= count[i];
			}
			*K = k;
			memcpy(C, c, sizeof(double)*3*k);
			memcpy(P, p, sizeof(int)*n);
			memcpy(Count, count, sizeof(int)*k);
			*err = variation;
			break;
		}
	}
}

bool split(const double*a, int n, double S0,int K0, int N0,
	int*K, double* C, int* P, int* Count)
{
	double* sigma = new double [(*K)*3];
	memset(sigma, 0, sizeof(double)*(*K)*3);
	double* d = new double [*K];
	memset(d, 0, sizeof(double)*(*K));

	for (int i = 0; i < n; i++)
	{
		sigma[P[i]*3] += (a[i*3]-C[P[i]*3])*(a[i*3]-C[P[i]*3]);
		sigma[P[i]*3+1] += (a[i*3+1]-C[P[i]*3+1])*(a[i*3+1]-C[P[i]*3+1]);
		sigma[P[i]*3+2] += (a[i*3+2]-C[P[i]*3+2])*(a[i*3+2]-C[P[i]*3+2]);
		d[P[i]] += dist(a, C, i, P[i]);
	}
	for (int i = 0; i < (*K)* 3; i++)
	{
		sigma[i] = sqrt(sigma[i]/Count[i/3]);
	}
	double D = 0;
	for (int i = 0; i < *K; i++)
	{	
		D += d[i];
		d[i] /= Count[i];
	}
	D /= n;
	double max = -INF;
	int t = -1;
	for (int i = 0; i < 3*(*K); i++)
	{
		if (sigma[i] > max)
		{
			max = sigma[i];
			t = i;
		}
	}
	int id = t/3;
	double* C2 = new double [6];
	memcpy(C2, C+3*id, sizeof(double)*3);
	memcpy(C2+3,C+3*id, sizeof(double)*3);
	C2[t-id*3] += 0.5*sigma[t];
	C2[t-id*3+3] -= 0.5*sigma[t];
	if ( (sigma[t]>S0)&&( (*K < K0/2)||(d[id]>D&&Count[id]>2*N0)))
	{
		memcpy(C+3*id, C2, sizeof(double)*3);
		memcpy(C+3*(*K), C2+3, sizeof(double)*3);
		*K = *K + 1;
		return true;
	}
	return false;
}

bool merge(const double*a, int n, double D0,
	int*K, double* C, int* P, int* Count)
{
	double* d = new double [(*K)*(*K)];
	memset(d, 0, sizeof(double)*(*K)*(*K));
	double min = INF;
	int I1 = -1;
	int J1 = -1;
	for (int i = 0; i < *K; i++)
	{
		d[i*(*K)+i] = INF;
		for (int j = 0; j < *K; j++)
		{
			if (i!=j) d[i*(*K)+j] = dist(C,C,i,j);
			if (d[i*(*K)+j] < min)
			{
				min = d[i*(*K)+j];
				I1 = i;
				J1 = j;
			}
		}
	}
	if ( min < D0 ) 
	{
		int I = I1 < J1 ? I1 : J1;
		int J = I1 < J1 ? J1 : I1;
		for(int i = 0; i < n; i++)
		{
			if (P[i]==J) P[i] = I;
		}
		C[3*I] = (C[3*I]*Count[I]+C[3*J]*Count[J])/(Count[I]+Count[J]);
		C[3*I+1] = (C[3*I+1]*Count[I]+C[3*J+1]*Count[J])/(Count[I]+Count[J]);
		C[3*I+2] = (C[3*I+2]*Count[I]+C[3*J+2]*Count[J])/(Count[I]+Count[J]);
		for(int i = J; i < *K-1; i++)
		{
			memcpy(C+3*J, C+3*(J+1), sizeof(double)*3);
			Count[i] = Count[i+1];
		}
		Count[I] += Count[J];
		*K = *K -1;
		return true;
	}
	return false;
}

void isodata(double*a, int n, int K0,
	int N0, double D0, double S0, int MAXSTEPS, 
	int* P, double* C, int *K)
{
	int k = K0;
	int* p = new int [n];
	int* count = new int [MAXK];
	double* c = new double [MAXK];
	memset(p, 0, sizeof(int)*n);
	memset(c, 0,sizeof(double)*MAXK);
	memset(count, 0, sizeof(int)*MAXK);
	bool flag = false;
	int err = 0;
	int step = 0;
	initCast(a,n,K0,c,p,count); 
	while (!flag)
	{
		step++;
		while(1)
		{
			cast3(a,n,N0,&k,c,p,count,&err);
			if (step % 2 == 0 ) break;
			if (!split(a,n,S0,K0,N0,&k,c,p,count)) break;
		}
		cout<<err<<endl;
		if ( err == 0 || step == MAXSTEPS) flag = true;
		merge(a,n,D0,&k,c,p,count);
	}
	memcpy(P, p, sizeof(int)*n);
	memcpy(C, c, sizeof(double)*3*k);
	*K = k;
}


void mexFunction(int nlhs, mxArray* plhs[], int nrhs, mxArray* prhs[])
{
	if ( nrhs != 6 )
		{	
			mexPrintf("\n1 data 3*n");
			mexPrintf("\n2 initial cluster number");
			mexPrintf("\n3 min number of birds in a cluster");
			mexPrintf("\n4 max merge distance");
			mexPrintf("\n5 min separate sigma");
			mexPrintf("\n6 max number of iteration\n");
			mexErrMsgTxt("6 input arguments are required.");
		}
	int m = mxGetM(prhs[0]);
	int n = mxGetN(prhs[0]);
	if (m != 3)
		mexErrMsgTxt("First dimension must be 3");
	
	double* a = mxGetPr(prhs[0]);
	int K0 = (int)mxGetScalar(prhs[1]);
	int N0 = (int)mxGetScalar(prhs[2]);
	double D0 = (double)mxGetScalar(prhs[3]);
	double S0 = (double)mxGetScalar(prhs[4]);
	int MAXSTEPS = (int)mxGetScalar(prhs[5]);
	
	int* P = new int [n];
	double* C = new double [MAXK*3];
	memset(P, 0, sizeof(int)*n);
	memset(C, 0, sizeof(double)*MAXK*3);
	int K = 0;
	isodata(a,n,K0,N0,D0,S0,MAXSTEPS,P,C,&K);
	
	const int N = mxGetNumberOfDimensions(prhs[0]);
	const int* D = mxGetDimensions(prhs[0]);
	int* DIM = new int [N-1];
	memcpy(DIM, D+1, sizeof(int)*(N-1));
	plhs[0] = mxCreateNumericMatrix(1, 1, mxDOUBLE_CLASS, mxREAL);
	plhs[1] = mxCreateNumericMatrix(3, K, mxDOUBLE_CLASS, mxREAL);
	plhs[2] = mxCreateNumericArray(N - 1, DIM, mxINT32_CLASS, mxREAL);
	
	double* mxK = mxGetPr(plhs[0]);
	double* mxC = mxGetPr(plhs[1]);
	double* mxP = mxGetPr(plhs[2]);
	*mxK = K;
	memcpy(mxC, C, sizeof(double)*3*K);
	memcpy(mxP, P, sizeof(int)*n);
}












































